import numpy as np
import albumentations as A
import plotly.graph_objects as go
from plotly.subplots import make_subplots
def read_obj(filename):
triangles = []
vertices = []
with open(filename) as file:
for line in file:
components = line.strip(' \n').split(' ')
if components[0] == "f": # face data
# e.g. "f 1/1/1/ 2/2/2 3/3/3 4/4/4 ..."
indices = list(
map(lambda c: int(c.split('/')[0]) - 1, components[1:]))
for i in range(0, len(indices) - 2):
triangles.append(indices[i:i + 3])
elif components[0] == "v": # vertex data
# e.g. "v 30.2180 89.5757 -76.8089"
vertex = list(map(lambda c: float(c), components[1:]))
vertices.append(vertex)
return np.array(vertices), np.array(triangles)
def show_augmentation(original, augmented):
fig = make_subplots(rows=1,
cols=2,
specs=[[{
'type': 'scatter3d'
}, {
'type': 'scatter3d'
}]])
before = go.Scatter3d(x=original[:, 0],
y=original[:, 1],
z=original[:, 2],
mode="markers",
marker=dict(size=2,
color=original[:, 3]))
after = go.Scatter3d(x=augmented[:, 0],
y=augmented[:, 1],
z=augmented[:, 2],
mode="markers",
marker=dict(size=2,
color=augmented[:, 3]))
fig.add_trace(before, row=1, col=1)
fig.add_trace(after, row=1, col=2)
fig.show()
teapot, _ = read_obj("images/teapot.obj")
teapot = np.hstack((teapot, np.zeros((len(teapot), 1))))
teapot[:,[0, 1, 2]] = teapot[:,[0, 2, 1]]
teapot[teapot[:, 0] > 2, 3] = 1
teapot[teapot[:, 0] < -2, 3] = 2
teapot[teapot[:, 2] > 2.5, 3] = 3
augmentation = A.Scale3d(scale_limit=[0.1, 0.1, 0.1], bias=[1.5, 1, 1], always_apply=True)
augmented_teapot = augmentation(points=teapot.copy())["points"]
show_augmentation(teapot, augmented_teapot)
augmentation = A.RotateAroundAxis3d(axis=[0, 1, 0], rotation_limit=np.pi/3, always_apply=True)
augmented_teapot = augmentation(points=teapot.copy())["points"]
show_augmentation(teapot, augmented_teapot)
augmentation = A.Crop3d(x_min=0, x_max=3, y_min=0, y_max=2, z_min=0, z_max=3, always_apply=True)
augmented_teapot = augmentation(points=teapot.copy())["points"]
show_augmentation(teapot, augmented_teapot)
augmentation = A.RandomMove3d(x_min=5, x_max=7, y_min=10, y_max=20, z_min=0, z_max=3, always_apply=True)
augmented_teapot = augmentation(points=teapot.copy())["points"]
show_augmentation(teapot, augmented_teapot)
augmentation = A.Center3d(always_apply=True)
augmented_teapot = augmentation(points=augmented_teapot)["points"]
show_augmentation(teapot, augmented_teapot)
augmentation = A.Flip3d(axis=[0,0,1], always_apply=True)
augmented_teapot = augmentation(points=teapot.copy())["points"]
show_augmentation(teapot, augmented_teapot)
augmentation = A.RandomDropout3d(always_apply=True)
augmented_teapot = augmentation(points=teapot.copy())["points"]
show_augmentation(teapot, augmented_teapot)